from functools import cached_property
from pathlib import Path
import logging
from pydantic import BaseModel, Field
from pysam import TabixFile, VariantFile, VariantHeader, tabix_compress, tabix_index  # pylint: disable=E0611
from .cnv import CNV

logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
logger = logging.getLogger("Annotate CNV")


class Annotation(BaseModel):
    genepred_file: Path | None = Field(default=None, description="Representative NCBI RefSeq file")
    gene_id_file: Path | None = Field(default=None, description="Gene to Entrez ID file")
    cytoband_file: Path | None = Field(default=None, description="CytoBand file")
    bed_files: list[Path] | None = Field(default_factory=list, description="BED files")
    overlap: float = Field(default=0.7, description="Overlap rate of frequency annotation file")

    @cached_property
    def gene_ids(self) -> dict[str, str]:
        with open(self.gene_id_file) as reader:
            return dict(map(lambda x: (f"{x[0]}:{x[1]}", x[2]), map(lambda y: y.strip().split("\t"), reader)))

    @cached_property
    def annotation_header(self) -> VariantHeader:
        header = VariantHeader()
        if self.gene_id_file:
            header.info.add(id="GENE", number=".", type="String", description="Gene Symbol")
            header.info.add(id="GENE_ID", number=".", type="String", description="Gene Entrez ID")
            header.info.add(id="DETAIL", number=".", type="String", description="Gene Detail")
        if self.cytoband_file:
            header.info.add(id="LOCATION", number=1, type="String", description="Location")
        if self.bed_files:
            for bed_file in self.bed_files:
                with TabixFile(str(bed_file)) as reader:
                    line = next(filter(lambda x: x.startswith("##INFO"), reader.header), "")
                    if line:
                        header.add_line(line)
                    else:
                        name = bed_file.name.split(".", 1)[0]
                        header.info.add(id=name, number=1, type="String", description=name)
        return header

    @cached_property
    def genepred_tabix(self) -> TabixFile | None:
        return TabixFile(str(self.genepred_file)) if self.genepred_file else None

    @cached_property
    def cytoband_tabix(self) -> TabixFile | None:
        return TabixFile(str(self.cytoband_file)) if self.cytoband_file else None

    @cached_property
    def bed_tabixes(self) -> dict[str, TabixFile]:
        tabixes = {}
        if self.bed_files:
            for bed_file in self.bed_files:
                tabix = TabixFile(str(bed_file))
                name = bed_file.name.split(".", 1)[0]
                line = next(filter(lambda x: x.startswith("##INFO"), tabix.header), "")
                if line:
                    header = VariantHeader()
                    header.add_line(line)
                    name = header.info.keys()[0]
                tabixes.setdefault(name, tabix)
        return tabixes

    def do_annotation(self, input_file: Path, output_file: Path):
        logger.info(f"Annotate {input_file} to {output_file}")
        outfile = output_file.with_suffix("") if output_file.suffix == ".gz" else output_file
        with VariantFile(str(input_file)) as reader, open(outfile, "w") as writer:
            reader.header.merge(self.annotation_header)
            writer.write(str(reader.header))
            for record in reader:
                cnv = CNV.from_variant_record(record)
                if self.genepred_tabix:
                    cnv.anno_gene(self.genepred_tabix, self.gene_ids)
                if self.cytoband_tabix:
                    cnv.anno_cytoband(self.cytoband_tabix)
                for db_name, db_tabix in self.bed_tabixes.items():
                    overlap = self.overlap if db_name.lower() in ["dgv", "decipher", "clinvar", "local"] else -1
                    cnv.anno_bed(db_tabix, db_name, overlap)
                record.info.update(cnv.info)
                writer.write(str(record))
        if output_file.suffix == ".gz":
            tabix_compress(str(outfile), str(output_file), force=True)
            tabix_index(str(output_file), force=True, preset="vcf")
            output_file.unlink(missing_ok=True)
        for tabix in filter(None, [self.genepred_tabix, self.cytoband_tabix] + list(self.bed_tabixes.values())):
            tabix.close()


class Annotator(Annotation):
    input_file: Path = Field(description="Input VCF file")
    output_file: Path = Field(description="Output VCF file")

    def run(self):
        self.do_annotation(self.input_file, self.output_file)
